package net.sansi.v3correctserver.service.photorec.rec;

import cn.hutool.core.date.StopWatch;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import net.sansi.v3correctserver.entity.IntPoint;
import net.sansi.v3correctserver.entity.color.PhotoColor;
import net.sansi.v3correctserver.entity.image.RecCmmImage;
import net.sansi.v3correctserver.entity.image.RecImage;
import net.sansi.v3correctserver.entity.photorec.LedNode;
import net.sansi.v3correctserver.exception.ServiceException;
import net.sansi.v3correctserver.util.CorrectUtil;
import net.sansi.v3correctserver.util.Matrix;
import net.sansi.v3correctserver.util.MonotoneUtil;
import net.sansi.v3correctserver.util.SortUtil;

import java.awt.*;
import java.nio.file.Path;
import java.util.List;
import java.util.Queue;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Created by Fang
 * create time : 2024/11/4
 * IDE : IntelliJ IDEA
 * <p>
 * 识别矩形屏幕照片基类
 */

@Slf4j
@Getter
public class RecRectScreen extends RecBase {

    // 方向数组
    private final int[] DX = new int[]{0, 1, 0, -1};
    private final int[] DY = new int[]{1, 0, -1, 0};


    // 识别到的所有灯点
    private List<LedNode> leds;
    // 构建出的矩阵
    private Matrix<LedNode> matrix;
    // 灯点纵横平均距离
    private int ledDistanceX, ledDistanceY;
    // 冲突点集合
    private Set<LedNode> conflictNodes;
    // 空点集合
    private Set<IntPoint> empties;
    // 所有灯点相加的亮度
    private double fullBrightness = 0;

    /**
     * 构造函数
     *
     * @param ledPhoto   led图片
     * @param blackPhoto 黑灯图片
     * @param fence      边界围栏
     */
    public RecRectScreen(RecImage ledPhoto, RecImage blackPhoto, List<IntPoint> fence) {
        super(ledPhoto, blackPhoto, fence);
    }


    /**
     * 开始识别
     *
     * @return this
     */
    public RecRectScreen toRec() {
        StopWatch watch = new StopWatch("图像识别");
        log.info("开始图像识别 ----------- ");
        watch.start("基准亮度计算");
        baseBright = computeBaseBright();
        watch.stop();
        watch.start("搜索灯点");
        leds = searchLeds();
        watch.stop();
        log.info("[searched leds count] {}", leds.size());
        watch.start("构建矩阵");
        matrix = buildMatrix(leds);
        watch.stop();
        watch.start("裁剪周边毛刺");
        clipRound();
        watch.stop();
        watch.start("计算灯点亮度");
        computeLedBright();
        watch.stop();
        log.info("完成图像识别 ------------ [{}]", watch.getTotal(TimeUnit.MILLISECONDS));
        //log.info("{}", watch.prettyPrint(TimeUnit.MILLISECONDS));
        return this;
    }

    /**
     * 裁剪周边无误识别的情况
     */
    private void clipRound() {
        int w = getMatrix().getWidth(), h = getMatrix().getHeight();
        int colThr = (int) Math.round(w * 0.02 + 1);
        int rowThr = (int) Math.round(h * 0.02 + 1);
        IntPoint lt = new IntPoint(0, 0), rb = new IntPoint(w - 1, h - 1);
        List<LedNode> top = matrix.filterRow(0, (x, y, node) -> node != null);
        if (top.size() <= colThr) lt.setY(1);
        List<LedNode> bottom = matrix.filterRow(h - 1, (x, y, node) -> node != null);
        if (bottom.size() <= colThr) rb.setY(h - 2);
        List<LedNode> right = matrix.filterCol(w - 1, (x, y, node) -> node != null);
        if (right.size() <= rowThr) rb.setX(w - 2);
        List<LedNode> left = matrix.filterCol(0, (x, y, node) -> node != null);
        if (left.size() <= rowThr) lt.setX(1);
        if (rb.getY() - lt.getY() == h - 1 && rb.getX() - lt.getX() == w - 1) return;
        matrix = matrix.clip(lt, rb);
        matrix.foreach((x, y, led) -> {
            if (led == null) return;
            led.setMatrixPos(new IntPoint(x, y));
        });
        clipRound();
    }

    /**
     * 计算灯点亮度
     */
    private void computeLedBright() {
        int w = ledDistanceX * 9 / 10 / 2, h = ledDistanceY * 9 / 10 / 2;
        matrix.foreach((mx, my, node) -> {
            if (node == null) return;
            IntPoint pos = node.getPhotoPos();
            int[] fence = new int[]{pos.getY() - h, pos.getX() + w, pos.getY() + h, pos.getX() - w};
            for (int x = fence[3]; x <= fence[1]; x++) {
                for (int y = fence[0]; y <= fence[2]; y++) {
                    PhotoColor pixelColor = getPixelColor(x, y);
                    node.addBright(pixelColor);
                    node.brightness += getPixelBrightness(x, y);
                }
            }
            node.brightFence = fence;
            fullBrightness += node.getBrightness();
        });
    }

    /**
     * 获取单色灯
     *
     * @param pos 灯点位置
     * @return 灯点
     */
    public LedNode getLedByBrightSort(double pos) {
        log.info("数量 {}", getMatrix().getItemSize());
        MonotoneUtil<LedNode> mon = new MonotoneUtil<>(
                (int) (pos * getMatrix().getItemSize()),
                (x, y) -> (int) (x.getBrightness() - y.getBrightness())
        );
        getMatrix().foreach((x, y, led) -> {
            if (led == null) return;
            mon.add(led);
        });
        return mon.get();
    }

    /**
     * 构建矩阵
     *
     * @param leds 灯点集合
     * @return 矩阵
     */
    private Matrix<LedNode> buildMatrix(List<LedNode> leds) {
        buildLedRelations(leds);
        int idx = 0, n = leds.size();
        Matrix<LedNode> tm = null;
        while (Math.abs(idx) >= 0 && Math.abs(idx) < n) {
            int i = idx >= 0 ? idx : n + idx;
            tm = buildMatrix(leds.get(i), leds);
            log.info("[build matrix] [idx] {} [w] {} [h] {} [pixels] {} [empty] {} [conflict] {}",
                    idx,
                    tm.getWidth(),
                    tm.getHeight(),
                    tm.getItemSize(),
                    empties.size(),
                    conflictNodes.size()
            );
            if ((tm.getItemSize() << 1) > n) return tm;
            idx = idx >= 0 ? -(idx + 1) : -(idx * 2);
        }
        throw new ServiceException("构建矩阵失败！");
    }

    /**
     * 构建矩阵
     *
     * @param start 起点
     * @param leds  灯点集合
     * @return 构建出的矩阵
     */
    private Matrix<LedNode> buildMatrix(LedNode start, List<LedNode> leds) {
        leds.forEach(node -> node.setMatrixPos(null));
        conflictNodes = new HashSet<>();
        empties = new HashSet<>();
        Matrix<LedNode> tMatrix = new Matrix<>(getLedPhoto().getWidth() / 4, getLedPhoto().getHeight() / 4);
        Queue<LedNode> queue = new LinkedList<>();
        Set<LedNode> vised = new HashSet<>();
        IntPoint startPos = new IntPoint(tMatrix.getWidth() >> 1, tMatrix.getHeight() >> 1);
        start.setMatrixPos(startPos);
        queue.add(start);
        int[] xValidDistance = new int[]{ledDistanceX * 2 / 3, ledDistanceX * 4 / 3};
        int[] yValidaDistance = new int[]{ledDistanceY * 2 / 3, ledDistanceY * 4 / 3};
        int[][] validDistance = new int[][]{yValidaDistance, xValidDistance, yValidaDistance, xValidDistance};
        int[] fence = new int[]{startPos.getY(), startPos.getX(), startPos.getY(), startPos.getX()};
        while (!queue.isEmpty()) {
            LedNode node = queue.poll();
            if (vised.contains(node)) continue;
            vised.add(node);
            IntPoint matrixPos = node.getMatrixPos();
            int mx = matrixPos.getX(), my = matrixPos.getY();

            LedNode orgNode = tMatrix.get(mx, my);
            if (orgNode != null) {
                if (orgNode != node) {
                    conflictNodes.add(orgNode);
                    conflictNodes.add(node);
                }
                continue;
            }

            tMatrix.set(mx, my, node);
            fence[0] = Math.min(fence[0], my);
            fence[1] = Math.max(fence[1], mx);
            fence[2] = Math.max(fence[2], my);
            fence[3] = Math.min(fence[3], mx);

            for (int dir = 0; dir < 4; dir++) {
                LedNode nextNode = node.neighbors[dir];
                if (nextNode == null) continue;
                int dis = computeLedDistance(node.getPhotoPos(), nextNode.getPhotoPos(), dir);
                int nx = mx + DX[dir], ny = my + DY[dir];
                if (dis > validDistance[dir][1]) continue;
                if (dis < validDistance[dir][0]) {
                    node.setMatrixPos(new IntPoint(mx, my));
                    conflictNodes.add(node);
                    conflictNodes.add(nextNode);
                } else if (!vised.contains(nextNode)) {
                    nextNode.setMatrixPos(new IntPoint(nx, ny));
                    queue.add(nextNode);
                }
            }
        }

        int mw = fence[1] - fence[3] + 1, mh = fence[2] - fence[0] + 1;
        Matrix<LedNode> ansMatrix = new Matrix<>(mw, mh);
        for (int x = fence[3]; x <= fence[1]; x++) {
            for (int y = fence[0]; y <= fence[2]; y++) {
                int ax = x - fence[3], ay = y - fence[0];
                LedNode node = tMatrix.get(x, y);
                if (node == null) empties.add(new IntPoint(ax, ay));
                else if (!conflictNodes.contains(node)) {
                    node.setMatrixPos(new IntPoint(ax, ay));
                    ansMatrix.set(ax, ay, node);
                }
            }
        }
        return ansMatrix;
    }

    private void buildLedRelations(List<LedNode> leds) {
        Matrix<LedNode> borderMatrix = new Matrix<>(getLedPhoto().getWidth(), getLedPhoto().getHeight());
        SortUtil<Integer> sortYDis = new SortUtil<>(Comparator.comparingInt(a -> a));
        SortUtil<Integer> sortXDis = new SortUtil<>(Comparator.comparingInt(a -> a));
        SortUtil<Integer>[] distances = new SortUtil[]{sortYDis, sortXDis};
        for (LedNode led : leds)
            led.borderPixels.forEach(pixel -> {
                int[] p = CorrectUtil.long2Xy(pixel);
                borderMatrix.set(p[0], p[1], led);
            });
        for (LedNode thsLed : leds) {
            IntPoint letPos = thsLed.getPhotoPos();
            for (int dir = 0; dir < 4; dir++) {
                int nx = letPos.getX() + DX[dir], ny = letPos.getY() + DY[dir];
                while (isInFence(nx, ny)) {
                    LedNode neighbor = borderMatrix.get(nx, ny);
                    if (neighbor != null && neighbor != thsLed) {
                        thsLed.neighbors[dir] = neighbor;
                        if (dir < 2) {
                            distances[dir].add(computeLedDistance(thsLed.getPhotoPos(), neighbor.getPhotoPos(), dir));
                        }
                        break;
                    }
                    nx += DX[dir];
                    ny += DY[dir];
                }
            }
        }
        ledDistanceX = sortXDis.get(0.5, 0);
        ledDistanceY = sortYDis.get(0.5, 0);
        log.info("[led distance x] {} [led distance y] {}", ledDistanceX, ledDistanceY);
    }


    /**
     * 搜索相片中的所有灯点
     *
     * @return this
     */
    private List<LedNode> searchLeds() {
        int w = getLedPhoto().getWidth(), h = getLedPhoto().getHeight();
        int span = 3;
        boolean[][] vised = new boolean[w][h];
        List<LedNode> leds = new ArrayList<>();
        for (int x = 0; x < w; x += span) {
            for (int y = 0; y < h; y += span) {
                if (!vised[x][y] && getPixelBrightness(x, y) >= baseBright && isInFence(x, y)) {
                    leds.add(bfsLed(x, y, vised));
                }
            }
        }
        leds.sort(Comparator.comparingInt(LedNode::getArea));
        int medianArea = leds.get(leds.size() >> 1).getArea();
        return leds.stream().filter(node -> {
            int area = node.getArea();
            return area > medianArea * 0.2 && area < medianArea * 2;
        }).toList();
    }

    /**
     * BFS 搜索所有灯点像素
     *
     * @param x     开始x坐标
     * @param y     开始y坐标
     * @param vised 是否访问
     * @return 该灯点
     */
    private LedNode bfsLed(int x, int y, boolean[][] vised) {
        LedNode resNode = new LedNode();
        Queue<int[]> queue = new LinkedList<>();
        queue.add(new int[]{x, y});
        while (!queue.isEmpty()) {
            int[] cur = queue.poll();
            int tx = cur[0], ty = cur[1];
            if (vised[tx][ty]) continue;
            vised[tx][ty] = true;
            boolean isBorder = false;
            for (int dir = 0; dir < 4; dir++) {
                int nx = tx + DX[dir], ny = ty + DY[dir];
                if (!isInFence(nx, ny) || getPixelBrightness(nx, ny) < baseBright) {
                    isBorder = true;
                    continue;
                }
                if (vised[nx][ny]) continue;
                queue.add(new int[]{nx, ny});
            }
            if (isBorder) resNode.addPixel(tx, ty, isBorder);
        }
        int left = resNode.ledFence[3], right = resNode.ledFence[1];
        while (left <= right) {
            int top = resNode.ledFence[0], bottom = resNode.ledFence[2];
            while (!vised[left][top]) top++;
            while (!vised[left][bottom]) bottom--;
            while (top < bottom) {
                //resNode.addPixel(left, top, false);
                //resNode.addPixel(left, bottom, false);
                top++;
                bottom--;
            }
            left++;
        }
        return resNode;
    }


    /**
     * 标记边界像素
     *
     * @return this
     */
    public RecRectScreen markBorders() {
        for (LedNode led : leds) {
            markLed(led, true, false);
        }
        return this;
    }

    /**
     * 标记矩阵灯点
     *
     * @return this
     */
    public RecRectScreen markMatrix() {
        matrix.foreach((int x, int y, LedNode node) -> {
            if (node == null) return;
            markLed(node, true, true);
        });
        return this;
    }

    /**
     * 绘制一个灯点信息
     *
     * @param led          灯点
     * @param isDrawBorder 是否绘制轮廓
     * @param isDrawFence  是否绘制亮度围栏
     */
    private void markLed(LedNode led, boolean isDrawBorder, boolean isDrawFence) {
        RecImage ledImage = getLedPhoto();
        if (isDrawBorder) {
            for (long p : led.borderPixels) {
                int[] pos = CorrectUtil.long2Xy(p);
                ledImage.setColor(pos[0], pos[1], new PhotoColor(Color.YELLOW));
            }
        }
        if (isDrawFence) {
            int[] fence = led.brightFence;
            ledImage.drawRect(
                    new IntPoint(fence[3], fence[0]),
                    new IntPoint(fence[1], fence[2]),
                    new PhotoColor(Color.CYAN)
            );
        }
    }

    /**
     * 保存识别结果矩阵图
     *
     * @param path 路径
     */
    public void saveMatrix(Path path) {
        RecCmmImage image = new RecCmmImage(matrix.getWidth(), matrix.getHeight());
        image.setBackgroundColor(new PhotoColor(Color.BLACK));
        matrix.foreach((x, y, node) -> {
            if (node == null) return;
            image.setColor(x, y, new PhotoColor(Color.CYAN));
        });
        image.save(path);
    }
}
